package com.lee.study;


import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.splitter.DocumentSplitters;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.store.embedding.EmbeddingMatch;
import dev.langchain4j.store.embedding.EmbeddingSearchRequest;
import dev.langchain4j.store.embedding.elasticsearch.ElasticsearchEmbeddingStore;
import org.apache.http.HttpHost;
import org.elasticsearch.client.RestClient;

import java.io.IOException;
import java.util.List;

public class ElasticsearchStore {
    private final RestClient restClient;
    private final ElasticsearchEmbeddingStore store;

    public ElasticsearchStore(String serverUrl, String index) {
        this.restClient = RestClient
                .builder(HttpHost.create(serverUrl))
                .build();
        this.store = ElasticsearchEmbeddingStore.builder()
                .restClient(restClient)
                .indexName(index)
                .build();

    }

    /**
     * 写入
     */
    public void ingest(List<Document> docs, EmbeddingService service) {
        var splitter = DocumentSplitters.recursive(500, 50);
        List<TextSegment> segments = docs.stream()
                .flatMap(d -> splitter.split(d).stream())
                .toList();
        List<Embedding> embeddings = service.embed(segments);
        store.addAll(embeddings, segments);
    }

    /**
     * 检索
     */
    public List<EmbeddingMatch<TextSegment>> search(String question, EmbeddingService service, int topK) {
        Embedding query = service.embed(question);
        EmbeddingSearchRequest req = EmbeddingSearchRequest.builder().queryEmbedding(query).maxResults(topK).build();
        return store.search(req).matches();
    }

    public void fini() {
        try {
            this.restClient.close();
        } catch (IOException e) {
            System.out.println("rest close failed, err:" + e);
        }
    }
}
